import numpy as np
import matplotlib.pyplot as plt
from optimizer import*
from collections import OrderedDict

def f(x,y):
    return x**2/20.0 + y**2

def df(x,y):
    return x/10.0,2.0*y

init_pos=(-7.0,2.0)#元组，所有优化器都从这个点开始寻优
params={}
params['x'],params['y']=init_pos[0],init_pos[1] #待优化的参数，存储了x和y的坐标
grads={}
grads['x'],grads['y']=0,0 #存储梯度的字典，后续会更新

#初始化每个最优化算法
optimizers=OrderedDict()
optimizers['SGD']=SGD(lr=0.95)
optimizers['Momentum']=Momentum(lr=0.1)
optimizers['AdaGrad']=AdaGrad(lr=1.5)
#还有一个 Adam 我还没有学会

idx=1 #子图索引，用于绘图布局

#遍历每一个要使用的优化器
for key in optimizers:
    optimizer=optimizers[key] #获取当前使用的优化器
    x_history=[] #记录x坐标的迭代历史，用于画轨迹
    y_history=[] #记录y坐标的迭代历史，用于画轨迹
    params['x'],params['y']=init_pos[0],init_pos[1] #重置参数到初始位置（每个优化器独立开始）

    #迭代30次（i 从 0 到 29）
    for i in range(30):
        #把当前的参数加入到列表，这样好画图
        x_history.append(params['x'])
        y_history.append(params['y'])

        grads['x'],grads['y']=df(params['x'],params['y']) #计算当前位置的梯度
        optimizer.update(params,grads) #用上面计算得到的梯度和当前使用的优化器来优化
    #绘制等高线
    x=np.arange(-10,10,0.01)#从-10到 10左闭右开，步长0.01，共 2000个点
    y=np.arange(-5,5,0.01)#从-5到 5 左闭右开，步长 0.01，共 1000个点

    X,Y=np.meshgrid(x,y) #每一个都根据另一个的大小扩展成二维数组
    Z=f(X,Y)#计算每个点的海拔高度

    mask=Z>7 #生成一个布尔矩阵：Z中大于 7 的位置为 True，否则为 False
    Z[mask]=0 #把Z中大于 7 的值设为 0，这样只显示函数值<=7的等高线，就不会显示远处等高线干扰了

    #绘制2行2列的子图，序号逐渐上升
    plt.subplot(2,2,idx)#
    idx+=1


    plt.plot(x_history,y_history,'o-',color='red')#画优化轨迹，圆点+红色实线
    plt.contour(X,Y,Z) #画目标函数等高线（默认是黑色线条，展示“地形”）
    plt.ylim(-10,10) #y范围
    plt.xlim(-10,10) #x范围
    plt.plot(0,0,'+') #标记位置
    plt.title(key) #子图标题
    plt.xlabel('x') #x轴标签
    plt.ylabel('y') #y轴标签
plt.show()

